from bzrc import BZRC, Command, ShootCommand
from PriorityQueue import PriorityQueue
import numpy 
import sys, math, time, optparse, random, pprint
import plot_grid, plotsearch
from Gnuplot import GnuplotProcess

KP = 1
KD = 2
BASE_PROB = .5
TRUE_THRESH = .99999
FALSE_THRESH = .00005
SKIP_NUM = 1
ROOT2 = math.sqrt(2)
NODENUM = 800 * 800
END_COUNT = NODENUM - 30
PRENUM = 20
CORNER1 =  numpy.ones((7, 7))
CORNER2 =  numpy.ones((7, 7))
CORNER3 =  numpy.ones((7, 7))
CORNER4 =  numpy.ones((7, 7))
CORNER5 =  numpy.ones((7, 7))
CORNER6 =  numpy.ones((7, 7))
CORNER7 =  numpy.ones((7, 7))
CORNER8 =  numpy.ones((7, 7))
FIRST_MATCH = 4 * 4
SECOND_MATCH = 7 * 7 - 3 * 3
VARIANCE = 1
TICKS_TO_SEARCH = 50

class Agent(object):

	def __init__(self, bzrc, numBots):
		self.bzrc = bzrc
		self.constants = self.bzrc.get_constants()
		self.commands = []
		#self.pfield = PField(self.bzrc)
		self.numBots = numBots
		self.lastError = [0 for _ in range(numBots)]
		self.lastTime = [time.time() for _ in range(numBots)]
		self.size = int(self.constants["worldsize"])
		self.truePos = float(self.constants['truepositive'])
		self.trueNeg = float(self.constants['truenegative'])
		self.probGrid = numpy.zeros( (self.size, self.size) )
		self.lastTime = numpy.ones(numBots)
		self.searchFlags = numpy.zeros(numBots)
		self.ticksSpent = numpy.zeros(numBots)
		self.myPath = {}
		self.probGrid.fill(BASE_PROB)
		plot_grid.init_window(self.size, self.size)
		plot_grid.update_grid(numpy.array(self.probGrid))
		self.search = UCS(self.size)
		self.done = False
		#self.gp = GnuplotProcess(persist=1)
		#self.gp.write(plotsearch.gnuplot_header(0, self.size))
		#self.gp.write(plotsearch.start_multiplot())
		self.tickCount = 0
		self.settledCount = 0
		#Create templates
		for i in range(0, 7):
			for j in range(0, 7):
				CORNER5[i,j] = -1
		for i in range(0, 4):
			for j in range(0, 4):
				CORNER5[i,j] = 1
		for i in range(0, 3):
			for j in range(0, 3):		
				CORNER1[i,j] = -1
		CORNER2 = numpy.rot90(CORNER1)
		CORNER3 = numpy.rot90(CORNER2)
		CORNER4 = numpy.rot90(CORNER3)
		CORNER6 = numpy.rot90(CORNER5)
		CORNER7 = numpy.rot90(CORNER6)
		CORNER8 = numpy.rot90(CORNER7)

	def use_occgrid(self, index):
		pos = self.pos[index]
		occgrid = self.occgrid[index]
		
		for x in range(0, occgrid.shape[0]):
			for y in range(0, occgrid.shape[1]):
				self.update_probs(x + pos[0] + self.size / 2 , y + pos[1] + self.size/2, occgrid[x, y])
				
	#Observation is true or false, true if an obstacle was observed
	def update_probs(self, x, y, observation):
		if self.probGrid[x, y] <= 0 or self.probGrid[x, y] >= 1:
			return

		probObs = 0 # Probability of the observation
		probGiven = 0 # Probability of observation given true state
		if observation:
			probObs = self.probGrid[x, y] * self.truePos + (1.0 - self.probGrid[x, y]) * (1.0 - self.trueNeg)
			probGiven = self.truePos
			self.probGrid[x, y] = self.probGrid[x, y] * probGiven / probObs
		else:
			probObs = (1.0 - self.probGrid[x, y]) * self.trueNeg + self.probGrid[x, y] * (1.0 -self.truePos)
			probGiven = self.trueNeg
			self.probGrid[x, y] = (1 - self.probGrid[x, y] * probGiven / probObs)
		if self.probGrid[x, y] < FALSE_THRESH:
			self.probGrid[x, y] = 0
		elif self.probGrid[x, y] > TRUE_THRESH:
			self.probGrid[x, y] = 1
		
		

	def tick(self):
		'''Some time has passed; decide what to do next'''
		# Get information from the BZRC server
		mytanks = self.bzrc.get_mytanks()
		pos, grids = self.bzrc.get_grids(self.numBots)
		self.mytanks = mytanks
		self.pos = pos
		self.occgrid = grids

		# Decide what to do with each of my tanks
		for tank in mytanks:
			# Reset my set of commands (we don't want to run old commands)
			self.commands = []
			self.done = self.behave(tank)
			# Send the commands to the server
			results = self.bzrc.do_commands(self.commands)
			if self.done:
				print "Finished searching space!"
				self.findCorners()
				break

		plot_grid.update_grid(numpy.array(self.probGrid))
		plot_grid.draw_grid()
					
		self.tickCount += 1

	def behave(self, bot):
		self.use_occgrid(bot.index)
		nodeNum = int(self.searchFlags[bot.index])
		
		if self.tickCount < PRENUM:
			target = (35 - self.size / 2, 150 * bot.index - self.size / 2 + 5)
		#elif self.tickCount < PRENUM * 2.5:
		#	target = (self.size / 2 - 35, 150 * bot.index - self.size / 2 + 5)
		elif nodeNum == 0:
			command = Command(bot.index, 0, 1)
			self.commands.append(command)
			results = self.bzrc.do_commands(self.commands)
			self.myPath[bot.index] = self.search.search(self.probGrid, (int(bot.x + self.size / 2), int(bot.y + self.size / 2)))
			path = self.myPath[bot.index]

			#self.gp.clear()
			#self.gp.write(plotsearch.plot_nodes(path, bot.index))
			#self.gp.flush()
			if not path:
				print "Finished searching!"
				return True
			elif len(path) > SKIP_NUM:
				target = (path[SKIP_NUM][0] - self.size / 2, path[SKIP_NUM][1] - self.size / 2)
				self.searchFlags[bot.index] = 1
			else:
				target = (path[-1][0] - self.size / 2, path[-1][1] - self.size / 2)
		else:
			path = self.myPath[bot.index]
			self.ticksSpent[bot.index] += 1
			if len(path) > SKIP_NUM + nodeNum and self.ticksSpent[bot.index] <= TICKS_TO_SEARCH:
				xPos = path[nodeNum][0] - self.size / 2
				yPos = path[nodeNum][1] - self.size / 2
				target = (xPos, yPos)
				dist = math.hypot(xPos - bot.x, yPos - bot.y)
				while dist < 35 or self.probGrid[xPos, yPos] == 0 or self.probGrid[target] == 1:
					self.searchFlags[bot.index] += SKIP_NUM
					nodeNum = int(self.searchFlags[bot.index])
					if nodeNum >= len(path):
						self.searchFlags[bot.index] = len(path) - 1
						break
					xPos = path[nodeNum][0] - self.size / 2
					yPos = path[nodeNum][1] - self.size / 2
					dist = math.hypot(xPos - bot.x, yPos - bot.y)
				target = (xPos, yPos)
			else:
				self.ticksSpent[bot.index] = 0
				target = (path[-1][0] - self.size / 2, path[-1][1] - self.size / 2)
				self.searchFlags[bot.index] = 0
							
		vel, ang = self.pd(bot, target)
		command = Command(bot.index, vel, ang)
		self.commands.append(command)
		return False

	def pd(self, bot, targetPos):
		targetAng = math.atan2(targetPos[1] - bot.y, targetPos[0] - bot.x)
		botAng = self.normalize_angle(bot.angle)
		error = self.normalize_angle(targetAng - botAng)
		now = time.time()
		delA = (error - self.lastError[bot.index]) / (now - self.lastTime[bot.index])
		a = KP * error + KD * delA
		self.lastError[bot.index] = error
		self.lastTime[bot.index] = now
		a = self.normalize_angle(a)
		return 1 - .75 * math.fabs(a / math.pi), a

	def normalize_angle(self, angle):
		'''Make any angle be between +/- pi.'''
		angle -= 2 * math.pi * int (angle / (2 * math.pi))
		if angle <= -math.pi:
			angle += 2 * math.pi
		elif angle > math.pi:
			angle -= 2 * math.pi
		return angle
		
	def findCorners(self):
		for i in range(4, self.size -4):
			for j in range(4, self.size -4):
				testArr = self.probGrid[i-3:i+4, j-3:j+4]
				#check if it matches a corner template
				if self.template_match(CORNER1, testArr, SECOND_MATCH) or self.template_match(CORNER2, testArr, SECOND_MATCH) or self.template_match(CORNER3, testArr, SECOND_MATCH) \
					or self.template_match(CORNER4, testArr, SECOND_MATCH) or self.template_match(CORNER5, testArr, FIRST_MATCH) or self.template_match(CORNER6, testArr, FIRST_MATCH) \
					or self.template_match(CORNER7, testArr, FIRST_MATCH) or self.template_match(CORNER8, testArr, FIRST_MATCH):
					print "Corner: ", (i - self.size / 2, j  - self.size / 2)
					
	def template_match(self, template, array, checkValue):
		totalSum = 0
		for i in range(0, 7):
			for j in range(0, 7):
				if array[i,j] < .5:
					array[i,j] = 0
				if array[i,j] > .5:
					array[i,j] = 1
				totalSum += template[i, j] * array[i,j]
		
				
		if (checkValue - VARIANCE) <= totalSum and totalSum <= (checkValue + VARIANCE):
			return True
		return False
					
		

class UCS:
	def __init__(self, size):
		self.done = False
		self.size = size
		
	def search(self, nodes, startPos):
		startTime = time.time()
		if self.done:
			return False
		self.nodes = nodes
		self.startPos = startPos
		self.lastNode = startPos
		self.frontier = PriorityQueue()
		self.frontier.enqueue(0, startPos, None, 0)
		self.path = []
		self.searched = {}
		self.newNodes = []
		self.penalties = nodes
		exists = self.iter_search()
		if not exists:
			self.done = True
			return
		finalPath = [exists]
		curNode = exists
		nextNode = self.searched[exists]
		while nextNode != None:
			finalPath.insert(0, nextNode)
			curNode = nextNode
			nextNode = self.searched[curNode]
		delTime = time.time() - startTime
		print "Done searching! Took ", delTime, " Length ", len(finalPath)
		return finalPath
		
	def iter_search(self):
		while self.frontier.has_item():
			answer = self.iter_step()
			if answer:
				return answer
		return False
		
	def iter_step(self):
		priority, cost, curNode, parent = self.frontier.dequeue()
		if not curNode in self.searched:
			self.searched[curNode] = parent
		pos = int(curNode[0]), int(curNode[1])
		if pos[0] >= self.size or pos[1] >= self.size or pos[0] < 0 or pos[1] < 0:
			return False
		if self.nodes[pos] == BASE_PROB:
			return curNode
		self.check_node((pos[0]+5, pos[1]-5), cost, curNode, True)
		self.check_node((pos[0]-5, pos[1]-5), cost, curNode, True)
		self.check_node((pos[0]+5, pos[1]+5), cost, curNode, True)
		self.check_node((pos[0]-5, pos[1]+5), cost, curNode, True)
		self.check_node((pos[0], pos[1]-5), cost, curNode)
		self.check_node((pos[0], pos[1]+5), cost, curNode)
		self.check_node((pos[0]+5, pos[1]), cost, curNode)
		self.check_node((pos[0]-5, pos[1]), cost, curNode)
		return False
		
	def check_node(self, node, cost, parent, corner=False):
		if node[0] >= self.size or node[1] >= self.size or node[0] < 0 or node[1] < 0:
			return False
		if self.nodes[node] < 1 and not node in self.searched:
			newCost = cost + self.calc_cost(node, parent, corner)
			if not self.frontier.get_data_priority(node):
				self.frontier.enqueue(newCost, node, parent, newCost)
			
	def calc_cost(self, node, parent, corner):
		return (corner and ROOT2) or 1

def normalize_angle(angle):
	'''Make any angle be between +/- pi.'''
	angle -= 2 * math.pi * int (angle / (2 * math.pi))
	if angle <= -math.pi:
		angle += 2 * math.pi
	elif angle > math.pi:
		angle -= 2 * math.pi
	return angle
def main():
	# Process CLI arguments.
	parser = optparse.OptionParser(usage = "%prog [options]",
								   version = "%prog 0.1")
	parser.add_option("-o", "--host", type="string",dest="host",
			metavar="HOST",default='localhost',
			help="hostname for the server")
	parser.add_option("-p", "--port", type="int",dest="port",
			metavar="PORT",default=3907,
			help="port number for the server")
	parser.add_option("-n", "--num", type="int",dest="num",
			metavar="NUM",default=10,
			help="number of bots for this lab")

	(options, args) = parser.parse_args()

	# Connect.
	#bzrc = BZRC(options.host, int(options.port), debug=True)
	bzrc = BZRC(options.host, int(options.port))

	agent = Agent(bzrc, options.num)

	# Run the agent
	try:
		while not agent.done:
			agent.tick()
	except KeyboardInterrupt:
		print "Exiting due to keyboard interrupt."
		bzrc.close()



if __name__ == '__main__':
	main()
